#!/usr/bin/env python
import numpy as np
try:
    from astropy.io import fits as pf
except ImportError:
    import pyfits as pf

def extractCountsSpec(infile, srcname, outfile):
    """
    Extract counts spectrum and integrated flux in the countsSpec.fits
    """
    fhand = pf.open(infile)
    try:
        counts = fhand['COUNTS_SPECTRA'].data.field(srcname)
    except KeyError:
        raise IOError("'%s' is not found in the %s!" % (srcname, infile))
    try:
        flux = fhand['FLUXES'].data.field(srcname)
    except KeyError:
        print "'%s' may be diffuse source?" % srcname
        flux = np.zeros_like(counts)
    emin = fhand['EBOUNDS'].data.field('E_MIN')
    emax = fhand['EBOUNDS'].data.field('E_MAX')
    fhand.close()

    outdata = np.c_[emin, emax, flux, counts]
    np.savetxt(outfile, outdata, delimiter='\t', header='Emin[MeV]\tEmax[MeV]\tFlux[ph cm^-2 s^-1]\tCounts[ph]')

def cli():
    import argparse

    parser = argparse.ArgumentParser(description='Extract counts spectrum and integrated flux in the countsSpec.fits')
    parser.add_argument("coutsspec", type=str, help='input countsSpec.fits')
    parser.add_argument("srcname", type=str, help='the name of the source you want')
    parser.add_argument("-outfile", type=str, default='cspect_extract.dat', help='the name of output file')
    args = parser.parse_args()

    extractCountsSpec(infile=args.coutsspec, srcname=args.srcname, outfile=args.outfile)

if __name__ == '__main__': cli()
